#include "yolov3_clip_detector.h"
#include <algorithm>
#include "acl/acl.h"
#include <dirent.h>
#include <math.h>
#include <cstdlib>
#include <chrono>
#include <unistd.h>

AclLiteModel::AclLiteModel(const std::string modelPath, const size_t deviceId):modelPath(modelPath),deviceId(deviceId){
    std::cout<<"-- Creat instence of Resource AclLiteModel"<<std::endl;
}

AclLiteModel::~AclLiteModel(){
    aclError ret;
    // if (modelDesc != nullptr) {
    //     ret = aclmdlDestroyDesc(modelDesc);
    //     if (ret != ACL_SUCCESS) {
    //         std::cout<<"Unload model error: "<< ret<<std::endl;
    //     }
    //     modelDesc = nullptr;
    // }
    std::cout<<"start ~AclLiteModel"<<std::endl;
    ret = aclmdlUnload(modelId);
    if (ret != ACL_SUCCESS) {
        std::cout<<"Unload model error: "<< ret<<std::endl;
    }

    
    for (size_t i = 0; i < aclmdlGetDatasetNumBuffers(input); ++i) {
        aclDataBuffer* dataBuffer = aclmdlGetDatasetBuffer(input, i);
        aclDestroyDataBuffer(dataBuffer);
        dataBuffer = nullptr;
    }
    aclmdlDestroyDataset(input);
    input = nullptr;
    std::cout<<"aclmdlDestroy input success"<<std::endl;

    for (size_t i = 0; i < aclmdlGetDatasetNumBuffers(output); ++i) {
        aclDataBuffer* dataBuffer = aclmdlGetDatasetBuffer(output, i);
        void* data = aclGetDataBufferAddr(dataBuffer);
        (void)aclrtFree(data);
        (void)aclDestroyDataBuffer(dataBuffer);
        dataBuffer = nullptr;
    }

    (void)aclmdlDestroyDataset(output);
    output = nullptr;
    std::cout<<"aclmdlDestroy output success"<<std::endl;

    std::cout<<"outputHostLists size: "<<outputHostLists.size()<<", outputDeviceLists size: "<<outputDeviceLists.size()<<std::endl;
    for(uint32_t i=0; i<outinput_nums; i++){
        aclrtFreeHost(outputHostLists[i]);
        aclrtFreeHost(outputDeviceLists[i]);
    }
    
    std::cout<<"modelInputs size: "<<modelInputs.size()<<std::endl;
    for(uint32_t i=0; i<input_nums; i++){
        aclrtFree(modelInputs[i]);
    }

    ret = aclrtDestroyContext(context);
    if (ret != ACL_SUCCESS) {
        std::cout<<"destroy context failed, errorCode is : "<< ret<<std::endl;
    }
    context = nullptr;

    ret = aclrtResetDevice(deviceId);
    if (ret != ACL_SUCCESS) {
        std::cout<<"reset device failed, errorCode is : "<< ret<<std::endl;
    }
    std::cout<<"Reset device "<< deviceId<<std::endl; 
}


bool AclLiteModel::init()
{
    // open device
    aclError ret = aclrtSetDevice(deviceId);
    if (ret != ACL_SUCCESS) {
        std::cout<<"Acl open device "<<deviceId<<" failed, errorCode is : "<<ret<<std::endl;
        return false;
    }
    std::cout<<"Open device "<<deviceId<<" ok"<<std::endl;

    ret = aclrtCreateContext(&context, deviceId);
    if (ret != ACL_SUCCESS) {
        std::cout<<"Create acl context failed, errorCode is : "<<ret<<std::endl;
        return false;
    }
    
    // Gets whether the current application is running on host or Device
    ret = aclrtGetRunMode(& runMode);
    if (ret != ACL_SUCCESS) {
        std::cout<<"acl get run mode failed"<<std::endl;
        return false;
    }

    ret = aclmdlLoadFromFile(modelPath.c_str(), &modelId);
    if (ret != ACL_SUCCESS) {
        std::cout<<"Load model "<<modelPath<<" from file return "<<ret<<std::endl;
        return false;
    }
    std::cout<<"modelId: "<<modelId<<std::endl;
    modelDesc = aclmdlCreateDesc();
    if (modelDesc == nullptr) {
        std::cout<<"create model description failed"<<std::endl;
        return false;
    }

    ret = aclmdlGetDesc(modelDesc, modelId);
    if (ret != ACL_SUCCESS) {
        std::cout<<"get model description failed, modelId is "<<modelId<< ", errorCode is "<<static_cast<int32_t>(ret)<<std::endl;
        return false;
    }

    std::cout<<"create model description success"<<std::endl;
    // cteate output
    output = aclmdlCreateDataset();
    if (output == nullptr) {
        std::cout<<"can't create dataset, create output false"<<std::endl;
        return false;
    }

    outinput_nums = aclmdlGetNumOutputs(modelDesc);
    std::cout<<"outinput_nums: "<<outinput_nums<<std::endl;
    for (size_t i = 0; i < outinput_nums; ++i) {
        const char* name = aclmdlGetOutputNameByIndex(modelDesc, i);
        std::cout<<"output op name: "<<name<<std::endl;
        size_t output_size = aclmdlGetOutputSizeByIndex(modelDesc, i);

        void * outputDevice;
        ret = aclrtMalloc(&outputDevice, output_size, ACL_MEM_MALLOC_NORMAL_ONLY);
        if (ret != ACL_SUCCESS) {
            std::cout<<"can't malloc buffer, size is"<<output_size<<", create output failed"<<std::endl;
            return false;
        }

        aclDataBuffer* outputData = aclCreateDataBuffer(outputDevice, output_size);
        if (ret != ACL_SUCCESS) {
            std::cout<<"can't create data buffer, create output failed"<<std::endl;
            aclrtFree(outputDevice);
            return false;
        }

        ret = aclmdlAddDatasetBuffer(output, outputData);
        if (ret != ACL_SUCCESS) {
            std::cout<<"can't add data buffer, create output failed"<<std::endl;
            aclrtFree(outputDevice);
            aclDestroyDataBuffer(outputData);
            return false;
        }

        void * outputHost;
        ret = aclrtMallocHost(&outputHost, output_size);
        if (ret != ACL_SUCCESS) {
            std::cout<<"Create model input failed"<<std::endl;
            return false;
        }
        outputHostLists.push_back(outputHost);
        outputDeviceLists.push_back(outputDevice);
        outputSizeLists.push_back(output_size);
    }

    std::cout<<"create model output success"<<std::endl;

    std::cout<<"create model INPUT"<<std::endl;
    input = aclmdlCreateDataset();
    if (input == nullptr) {
        std::cout<<"Create input failed for create dataset failed"<<std::endl;
        return false;
    }

    input_nums = aclmdlGetNumInputs(modelDesc);
    std::cout<<"input_nums: "<<input_nums<<std::endl;
    if (input_nums == 0) {
        std::cout<<"Create input failed for no input data"<<std::endl;
        return false;
    }

    for (size_t i= 0; i<input_nums; i++){
        size_t InputSize = aclmdlGetInputSizeByIndex(modelDesc, i);
        std::cout<<"InputSize: "<<InputSize<<std::endl;
        
        void* modelInputBuf;
        aclError aclRet = aclrtMalloc(&modelInputBuf, (size_t)(InputSize), ACL_MEM_MALLOC_HUGE_FIRST);
        if (aclRet != ACL_SUCCESS) {
            std::cout<<"malloc device data buffer failed, aclRet is "<< aclRet<<std::endl;
            return false;
        }

        aclDataBuffer* dataBuf = aclCreateDataBuffer(modelInputBuf, InputSize);
        if (dataBuf == nullptr) {
            std::cout<<"Create data buffer error"<<std::endl;
            return false;
        }

        aclError ret = aclmdlAddDatasetBuffer(input, dataBuf);
        if (ret != ACL_SUCCESS) {
            std::cout<<"Add dataset buffer error "<<ret<<std::endl;
            ret = aclDestroyDataBuffer(dataBuf);
            if (ret != ACL_SUCCESS) {
                std::cout<<"Destroy dataset buffer error "<<ret<<std::endl;
            }
            dataBuf = nullptr;
            return false;
        }
        modelInputs.push_back(modelInputBuf);
        inputSizeLists.push_back(InputSize);
    }
    std::cout<<"create model input success"<<std::endl;
    return true;
}

aclrtContext AclLiteModel::getContext(){
    return context;
}

void AclLiteModel::Execute(std::vector<cv::Mat> &frames, std::vector<std::vector<float>> &outputs){
    aclrtSetCurrentContext(context);
    for (size_t i= 0; i<input_nums; i++){
        aclError ret = aclrtMemcpy(modelInputs[i], inputSizeLists[i], frames[i].ptr<uint8_t>(), inputSizeLists[i], ACL_MEMCPY_DEVICE_TO_DEVICE);
        if (ret != ACL_SUCCESS) {
            std::cout<<"Copy data to device failed, aclRet is "<<ret<<std::endl;
            break;
        }
    }
    
    aclError ret = aclmdlExecute(modelId, input, output);
    if (ret != ACL_SUCCESS) {
        std::cout<<"execute aclmdlExecute false "<<static_cast<int32_t>(ret)<<std::endl;
    }
    for (size_t i= 0; i<outinput_nums; i++){
        float * outFloatData0;
        int outputsize = outputSizeLists[i] /4;
        if (runMode == ACL_HOST) {
            ret = aclrtMemcpy(outputHostLists[i], outputSizeLists[i], outputDeviceLists[i], outputSizeLists[i], ACL_MEMCPY_DEVICE_TO_HOST);
            outFloatData0 = reinterpret_cast < float * > (outputHostLists[i]);
        }
        else
        {
            outFloatData0 = reinterpret_cast < float * > (outputDeviceLists[i]);            
        }
        std::vector<float> vec_0(outFloatData0, outFloatData0 + outputsize);
        outputs.push_back(vec_0);
    }                
  }


void AclLiteModel::Execute(std::vector<std::vector<float>> &features, std::vector<std::vector<float>> &outputs){
    aclrtSetCurrentContext(context);
    for (int i=0; i< input_nums; i++){
        aclError ret = aclrtMemcpy(modelInputs[i], inputSizeLists[i], features[i].data(), inputSizeLists[0], ACL_MEMCPY_DEVICE_TO_DEVICE);
        if (ret != ACL_SUCCESS) {
            std::cout<<"Copy data to device failed, aclRet is "<<ret<<std::endl;
            break;
        }
    }
    
    
    aclError ret = aclmdlExecute(modelId, input, output);
    if (ret != ACL_SUCCESS) {
        std::cout<<"execute aclmdlExecute false "<<static_cast<int32_t>(ret)<<std::endl;
    }
    for (size_t i= 0; i<outinput_nums; i++){
        float * outFloatData0;
        int outputsize = outputSizeLists[i] /4;
        if (runMode == ACL_HOST) {
            ret = aclrtMemcpy(outputHostLists[i], outputSizeLists[i], outputDeviceLists[i], outputSizeLists[i], ACL_MEMCPY_DEVICE_TO_HOST);
            outFloatData0 = reinterpret_cast < float * > (outputHostLists[i]);
        }
        else
        {
            outFloatData0 = reinterpret_cast < float * > (outputDeviceLists[i]);            
        }
        std::vector<float> vec_0(outFloatData0, outFloatData0 + outputsize);
        outputs.push_back(vec_0);
    }                
}

Yolov3CLIPDetector::Yolov3CLIPDetector(const std::string modelInfPath, const std::string modelVisPath, const size_t deviceIdInf, const size_t deviceIdVis):Model_inf(modelInfPath, deviceIdInf),Model_vis(modelVisPath, deviceIdVis)
{
  std::cout<<"-- Creat instence of Resource Yolov3CLIPDetector"<<std::endl;
}

void Yolov3CLIPDetector::DestroyResource()
{
    done = true;
    queueInf_pre.notify_all();
    queueVis_pre.notify_all();

    for (auto& th : threads) {
        if (th.joinable()) {
            th.join();
        }
    }

    // aclError ret = aclFinalize();
    // if (ret != ACL_SUCCESS) {
    //     std::cout<<"finalize acl failed, errorCode is : "<< ret<<std::endl;
    // }
    // std::cout<<"Finalize acl ok"<<std::endl;
}

Yolov3CLIPDetector::~Yolov3CLIPDetector()
{
  DestroyResource();
}


void Yolov3CLIPDetector::initThread(){
    for (u_int32_t i=0; i<preprocess_thread_num;i++){
        threads.emplace_back(&Yolov3CLIPDetector::modelInfExecute, this);
    }
    for (u_int32_t i=0; i<preprocess_thread_num;i++){
        threads.emplace_back(&Yolov3CLIPDetector::modelVisExecute, this);
    }
}

bool Yolov3CLIPDetector::init()
{
    aclError ret = aclInit("");
    if (ret != ACL_SUCCESS) {
        std::cout<<"Acl init failed, errorcode is: "<< ret<<std::endl;
        return false;
    }
    std::cout<<"Acl init ok"<<std::endl;

    bool flag = Model_inf.init();
    if (!flag){
        std::cout<<"init model inf false"<<std::endl;
        return false;
    }
    std::cout<<"init Model_inf success"<<std::endl;
    flag = Model_vis.init();
    if (!flag){
        std::cout<<"init model vis false"<<std::endl;
        return false;
    }
    std::cout<<"init Model_vis success"<<std::endl;
    initThread();
    return true;
}


void Yolov3CLIPDetector::modelInfExecute()
{
    message_pre premsg;
    message_result excutemsg;
    while(!done){
        std::vector<std::vector<float>> outputs;
        if (queueInf_pre.try_pop(premsg)){
            excutemsg.finalOutput.clear();
            std::vector<cv::Mat> frames = premsg.frames;
            {
                std::lock_guard<std::mutex> lock(outputInf_mutex);
                Model_inf.Execute(frames, outputs);
            }
            excutemsg.filenames = premsg.filenames;
            excutemsg.finalOutput = outputs;
            std::cout<<"modelInfExecute out: "<<outputs[0][0]<<std::endl;
            std::lock_guard<std::mutex> lock(resultsInf_mutex);
            resultsInf.push_back(excutemsg);
            if (--taskInf_counter == 0) {
                conda_Inf.notify_all();
            }
        }else{
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
    }
}

void Yolov3CLIPDetector::modelVisExecute()
{
    message_pre premsg;
    message_result excutemsg;
    while(!done){
        std::vector<std::vector<float>> outputs;
        if (queueVis_pre.try_pop(premsg)){
            excutemsg.finalOutput.clear();
            std::vector<cv::Mat> frames = premsg.frames;
            {
                std::lock_guard<std::mutex> lock(outputVis_mutex);
                Model_vis.Execute(frames, outputs);
            }
            excutemsg.filenames = premsg.filenames;
            excutemsg.finalOutput = outputs;
            std::lock_guard<std::mutex> lock(resultsVis_mutex);
            resultsVis.push_back(excutemsg);
            if (--taskVis_counter == 0) {
                conda_Vis.notify_all();
            }
        }else{
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
    }
}


void Yolov3CLIPDetector::modelAllExecute()
{
    message_pre premsg;
    message_result excutemsg;
    while(!done){
        std::vector<std::vector<float>> outputs;
        if (queueAll_pre.try_pop(premsg)){
            excutemsg.finalOutput.clear();
            std::vector<cv::Mat> frames = premsg.frames;
            {
                std::lock_guard<std::mutex> lock(outputVis_mutex);
                Model_all.Execute(frames, outputs);
            }
            excutemsg.filenames = premsg.filenames;
            excutemsg.finalOutput = outputs;
            std::lock_guard<std::mutex> lock(resultsVis_mutex);
            resultsAll.push_back(excutemsg);
            if (--taskAll_counter == 0) {
                conda_All.notify_all();
            }
        }else{
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
    }
}

void Yolov3CLIPDetector::wait_forInf_completion() {
    std::unique_lock<std::mutex> lock(resultsInf_mutex);
    conda_Inf.wait(lock, [this]{ return this->taskInf_counter == 0; });
}

void Yolov3CLIPDetector::wait_forVis_completion() {
    std::unique_lock<std::mutex> lock(resultsVis_mutex);
    conda_Vis.wait(lock, [this]{ return this->taskVis_counter == 0; });
}

void Yolov3CLIPDetector::wait_forAll_completion() {
    std::unique_lock<std::mutex> lock(resultsInf_mutex);
    conda_Inf.wait(lock, [this]{ return this->taskInf_counter == 0; });
}

std::vector<message_result> Yolov3CLIPDetector::inference(std::vector<std::string> &imageLists){
    {
        std::lock_guard<std::mutex> lock(resultsInf_mutex);
        resultsInf.clear();
        taskInf_counter = 0;
    }

    {
        std::lock_guard<std::mutex> lock(resultsVis_mutex);
        resultsVis.clear();
        taskVis_counter = 0;
    }
    
    for (int i=0; i<imageLists.size(); i++){
        message_pre premsg;
        cv::Mat img1 = cv::imread(imageLists[i]);
        cv::Mat reiszeInf_visible;
        cv::Mat reiszeVis_visible;
        cv::resize(img1, reiszeInf_visible, cv::Size(g_modelWidth_, g_modelHeight_),0, 0, cv::INTER_CUBIC);
        cv::resize(img1, reiszeVis_visible, cv::Size(g_modelWidth_, g_modelHeight_),0, 0, cv::INTER_CUBIC);
        premsg.filenames.push_back(imageLists[i]);
        premsg.frames.push_back(reiszeInf_visible);
        queueInf_pre.push(premsg);
        queueVis_pre.push(premsg);
        taskVis_counter++;
        taskInf_counter++;
    }

    wait_forInf_completion();
    wait_forVis_completion();



    std::vector<message_result> local_results;
    {
        std::lock_guard<std::mutex> lock(resultsInf_mutex);
        local_results = resultsInf;
    }

    return local_results;

}




